package crypto

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/x509"
	"errors"
	"fmt"
	"math/big"
	"strings"

	"github.com/tjfoc/gmsm/sm2"
	gmsmx509 "github.com/tjfoc/gmsm/x509"
)

const (
	//Sm4SymmetricPrivateKeySize sm4对称加密私钥大小
	Sm4SymmetricPrivateKeySize = 16
	//AesSymmetricPrivateKeySize aes对称加密私钥大小
	AesSymmetricPrivateKeySize = 32
)

type keyImporter interface {
	//KeyImport key import
	KeyImport(encoded []byte) (Key, error)
}
type ecdsaPublicKeyImporter struct {
	curve elliptic.Curve
}

type ecdsaPrivateKeyImporter struct {
	curve elliptic.Curve
}

//KeyImport key import
func (imp *ecdsaPublicKeyImporter) KeyImport(encoded []byte) (Key, error) {
	size := len(encoded)
	switch size {
	case 65: //加标志位的短公钥
		x, y := elliptic.Unmarshal(imp.curve, encoded)
		if x == nil {
			return nil, errors.New("invalid public key")
		}
		return &ecdsaPublicKey{
			&ecdsa.PublicKey{Curve: imp.curve, X: x, Y: y},
		}, nil
	case 64: //短公钥
		x := new(big.Int).SetBytes(encoded[0:32])
		y := new(big.Int).SetBytes(encoded[32:])
		return &ecdsaPublicKey{
			&ecdsa.PublicKey{Curve: imp.curve, X: x, Y: y},
		}, nil
	default:
		key, err := x509.ParsePKIXPublicKey(encoded)
		if err != nil {
			return nil, err
		}
		pubKey, ok := key.(*ecdsa.PublicKey)
		if ok {
			return &ecdsaPublicKey{
				&ecdsa.PublicKey{Curve: imp.curve, X: pubKey.X, Y: pubKey.Y},
			}, nil
		}
		return nil, errors.New("invalid public key")
	}

}

//KeyImport key import
func (imp *ecdsaPrivateKeyImporter) KeyImport(encoded []byte) (Key, error) {
	size := len(encoded)
	switch size {
	case 33:
		encoded = encoded[1:]
		fallthrough
	case 32: //短私钥
		priv := new(ecdsa.PrivateKey)
		priv.PublicKey.Curve = imp.curve

		priv.D = new(big.Int).SetBytes(encoded)
		N := imp.curve.Params().N
		// The priv.D must < N
		if priv.D.Cmp(N) >= 0 {
			return nil, fmt.Errorf("invalid private key, >=N")
		}
		// The priv.D must not be zero or negative.
		if priv.D.Sign() <= 0 {
			return nil, fmt.Errorf("invalid private key, zero or negative")
		}

		priv.PublicKey.X, priv.PublicKey.Y = priv.PublicKey.Curve.ScalarBaseMult(encoded)
		if priv.PublicKey.X == nil {
			return nil, errors.New("invalid private key")
		}
		return &ecdsaPrivateKey{priv}, nil
	case 97: //加标志位的短公钥
		encoded = encoded[1:]
		fallthrough
	case 96: //短公钥
		priv := new(ecdsa.PrivateKey)
		priv.PublicKey.Curve = imp.curve

		priv.D = new(big.Int).SetBytes(encoded[0:32])
		priv.PublicKey.X = new(big.Int).SetBytes(encoded[32:64])
		priv.PublicKey.Y = new(big.Int).SetBytes(encoded[64:])
		return &ecdsaPrivateKey{
			priv,
		}, nil
	default:
		priKey, err := x509.ParseECPrivateKey(encoded)
		if err != nil {
			return nil, err
		}

		return &ecdsaPrivateKey{
			priKey,
		}, nil

	}

}

//ParserPrivateKey 解析私钥
func ParserPrivateKey(encoded []byte) (PrivateKey, error) {
	return ParserPrivateKeyWithOpts(encoded, options.AsymmetricAlgorithm)
}

//ParserPrivateKeyWithOpts 使用指定算法解析私钥
func ParserPrivateKeyWithOpts(encoded []byte, algorithm string) (PrivateKey, error) {
	algorithm = strings.ToUpper(algorithm)
	var imp keyImporter
	switch algorithm {
	case "SM2":
		imp = &sm2PrivateKeyImporter{sm2.P256Sm2()}
	case "ECDSA":
		imp = &ecdsaPrivateKeyImporter{elliptic.P256()}
	default:
		return nil, fmt.Errorf(" No key importer registered for algorithm %s", algorithm)
	}
	k, err := imp.KeyImport(encoded)
	if err != nil {
		return nil, err
	}
	return k.(PrivateKey), nil
}

//ParserPublicKey 解析公钥
func ParserPublicKey(encoded []byte) (PublicKey, error) {
	return ParserPublicKeyWithOpts(encoded, options.AsymmetricAlgorithm)
}

//ParserPublicKeyWithOpts 解析公钥
func ParserPublicKeyWithOpts(encoded []byte, algorithm string) (PublicKey, error) {
	algorithm = strings.ToUpper(algorithm)
	var imp keyImporter
	switch algorithm {
	case "SM2":
		imp = &sm2PublicKeyImporter{sm2.P256Sm2()}
	case "ECDSA":
		imp = &ecdsaPublicKeyImporter{elliptic.P256()}
	default:
		return nil, fmt.Errorf(" No key importer registered for algorithm %s", algorithm)
	}
	k, err := imp.KeyImport(encoded)
	if err != nil {
		return nil, err
	}
	return k.(PublicKey), nil
}

//ParserPublicKeyWithSignatureAlgorithm 根据签名算法类型解析公钥
func ParserPublicKeyWithSignatureAlgorithm(encoded []byte, signatureAlgorithm int) (PublicKey, error) {
	//从签名算法获取非对称算法
	switch signatureAlgorithm {
	case 9, 10, 11, 12:
		return ParserPublicKeyWithOpts(encoded, "ECDSA")
	case 16, 17, 18:
		return ParserPublicKeyWithOpts(encoded, "SM2")
		//case 1, 2, 3, 4, 5, 6:
		//	return ParserPublicKeyWithOpts(encoded, "RSA")
		//case 7, 8:
		//	return ParserPublicKeyWithOpts(encoded, "DSA")
	default:
		return nil, fmt.Errorf(" Unknown Signature Algorithm %d", signatureAlgorithm)
	}
}

//ParserSymmetricKey 解析对称加密秘钥
func ParserSymmetricKey(encoded []byte) (Key, error) {
	if options.SymmetricAlgorithm == "SM4" {
		if encoded == nil || len(encoded) < Sm4SymmetricPrivateKeySize {
			return nil, errors.New("invalid Symmetric Key")
		}
		return &sm4PrivateKey{encoded[0:16], true}, nil
	}
	if len(encoded) != AesSymmetricPrivateKeySize {
		return nil, fmt.Errorf("invalid Key Length [%d]. Must be 32 bytes", len(encoded))
	}
	return &aesPrivateKey{encoded, true}, nil
}

//国密
type sm2PublicKeyImporter struct {
	curve elliptic.Curve
}

type sm2PrivateKeyImporter struct {
	curve elliptic.Curve
}

//KeyImport sm2 public key import
func (imp *sm2PublicKeyImporter) KeyImport(encoded []byte) (Key, error) {
	size := len(encoded)
	switch size {
	case 65: //加标志位的短公钥
		x, y := elliptic.Unmarshal(imp.curve, encoded)
		if x == nil {
			return nil, errors.New("invalid public key")
		}
		return &sm2PublicKey{
			&sm2.PublicKey{Curve: imp.curve, X: x, Y: y},
		}, nil
	case 64: //短公钥
		x := new(big.Int).SetBytes(encoded[0:32])
		y := new(big.Int).SetBytes(encoded[32:])
		return &sm2PublicKey{
			&sm2.PublicKey{Curve: imp.curve, X: x, Y: y},
		}, nil
	default:
		key, err := x509.ParsePKIXPublicKey(encoded)
		if err != nil {
			return nil, err
		}
		switch pubKey := key.(type) {
		case *ecdsa.PublicKey:
			return &sm2PublicKey{
				&sm2.PublicKey{Curve: imp.curve, X: pubKey.X, Y: pubKey.Y},
			}, nil
		case *sm2.PublicKey:
			return &sm2PublicKey{
				pubKey,
			}, nil
		default:
			return nil, errors.New("invalid public key")
		}
	}

}

//KeyImport sm2 private key import
func (imp *sm2PrivateKeyImporter) KeyImport(encoded []byte) (Key, error) {
	size := len(encoded)
	switch size {
	case 33:
		encoded = encoded[1:]
		fallthrough
	case 32: //短私钥
		priv := new(sm2.PrivateKey)
		priv.PublicKey.Curve = imp.curve

		priv.D = new(big.Int).SetBytes(encoded)
		N := imp.curve.Params().N
		// The priv.D must < N
		if priv.D.Cmp(N) >= 0 {
			return nil, fmt.Errorf("invalid private key, >=N")
		}
		// The priv.D must not be zero or negative.
		if priv.D.Sign() <= 0 {
			return nil, fmt.Errorf("invalid private key, zero or negative")
		}

		priv.PublicKey.X, priv.PublicKey.Y = priv.PublicKey.Curve.ScalarBaseMult(encoded)
		if priv.PublicKey.X == nil {
			return nil, errors.New("invalid private key")
		}
		return &sm2PrivateKey{priv}, nil
	case 97: //加标志位的短公钥
		encoded = encoded[1:]
		fallthrough
	case 96: //短公钥
		priv := new(sm2.PrivateKey)
		priv.PublicKey.Curve = imp.curve

		priv.D = new(big.Int).SetBytes(encoded[0:32])
		priv.PublicKey.X = new(big.Int).SetBytes(encoded[32:64])
		priv.PublicKey.Y = new(big.Int).SetBytes(encoded[64:])
		return &sm2PrivateKey{
			priv,
		}, nil
	default:
		priKey, err := gmsmx509.ParseSm2PrivateKey(encoded)
		if err != nil {
			return nil, err
		}
		priv := new(sm2.PrivateKey)
		priv.Curve = priKey.Curve
		priv.D = priKey.D
		priv.X, priv.Y = priKey.X, priKey.Y
		return &sm2PrivateKey{
			priv,
		}, nil

	}

}
